import pandas as pd
from os import path
from ARM_TGNCF.Parse_ARM_TGNCF import ARM_args, original_rating_name, cold_threshold
from loaddata.load_data_tool.load_original_events.ML_100K import load_ML_100K_Ratings
from loaddata.recoding_tool.df_tool import del_cold
from loaddata.recoding_tool.recoding import recoding_ratings


def load_events_data():
    """

    :return:
    """
    dataset_name = ARM_args.dataset

    path_dataset = path.join(path.dirname(__file__), '..', '..', 'Dataset', dataset_name)

    path_before_recoding = path.join(path_dataset, original_rating_name[dataset_name])

    path_recoding_ratings = path.join(path_dataset, 'ratings_recoding.dat')

    user_cold_threshold, item_cold_threshold = cold_threshold[dataset_name][:2]

    if not path.exists(path_recoding_ratings) or ARM_args.del_cold:
        if dataset_name == 'ML_100K':
            df_ratings = load_ML_100K_Ratings(path_before_recoding)
        else:
            df_ratings = None
            print('dataset is error')
            exit(1)

        if ARM_args.del_cold:
            df_ratings = del_cold(df_ratings, user_cold_threshold, 'userId')
            df_ratings = del_cold(df_ratings, item_cold_threshold, 'itemId')
            path_recoding_ratings = path.join(path_dataset, 'ratings_recode_cold.dat')
        df_ratings = recoding_ratings(path_recoding_ratings, df_ratings, sort_phase='userId')
        user_num = max(df_ratings['userId'])
        item_num = max(df_ratings['itemId'])
        print('user_num:', user_num, '  item_num', item_num)
        print('events num:', df_ratings.shape[0])

    else:
        df_ratings = pd.read_table(path_recoding_ratings,
                                   sep='\t',
                                   names=['userId', 'itemId', 'rating', 'timestamp'])

    df_ratings.sort_values(by='itemId', inplace=True, ascending=True)
    df_ratings.sort_values(by='userId', inplace=True, ascending=True)
    if ARM_args.events_is_ascend_time:
        df_ratings.sort_values(by='timestamp', inplace=True, ascending=True)

    return df_ratings[['userId', 'itemId', 'rating', 'timestamp']]


if __name__ == '__main__':
    print('dataset：', ARM_args.dataset)
    load_events_data()
